modelopt/torch/export/layer_utils.py (2)

1085-1102: ⚠️ Potential issue | 🟠 Major

Zero amax values are still collected into valid_amax_values, causing them to silently perpetuate when all experts are uncalibrated.

The new needs_amax condition (line 1142) correctly identifies all-zero amax tensors as invalid. However, the valid_amax_values collection loop (lines 1089–1098) only checks existing_amax is not None, so zero tensors are collected. When every expert has amax == 0:

1. valid_amax_values = [0, 0, ...]
2. target_amax = torch.max(stack([0, ...])) = 0
3. elif target_amax is None branch (line 1105) is skipped — weight-stat fallback never runs
4. needs_amax = True → quantizer.amax is set to 0 again
5. Warning misleadingly says "Setting it to 0.000000 (max from existing quantizers in current batch)"

The fix is to mirror the needs_amax predicate in the collection loop:

🐛 Proposed fix

     valid_amax_values = []
     for _, attr_name, quantizer in all_quantizers:
         existing_amax = getattr(quantizer, "amax", None)
-        if existing_amax is not None:
+        if existing_amax is not None and not (
+            isinstance(existing_amax, torch.Tensor) and torch.all(existing_amax == 0)
+        ):
             # Convert to tensor and add to collection
             if isinstance(existing_amax, torch.Tensor):
                 valid_amax_values.append(existing_amax.to(target_device))
             else:
                 valid_amax_values.append(
                     torch.tensor(existing_amax, dtype=torch.float32, device=target_device)
                 )

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/torch/export/layer_utils.py` around lines 1085 - 1102, When
collecting existing amax values into valid_amax_values (loop over all_quantizers
/ existing_amax), only append values that are non-zero using the same predicate
as needs_amax: convert existing_amax to a torch.Tensor on target_device first
and check it's not all zeros (e.g., tensor.ne(0).any()); skip appending if the
tensor is all zeros or None so target_amax won't be set to 0 when all experts
are uncalibrated, allowing the weight-stat fallback to run and avoiding
misleading warnings for quantizer.amax.

---

328-344: ⚠️ Potential issue | 🟠 Major

get_experts_list does not handle Qwen3.5 model type — will raise NotImplementedError at runtime.

is_moe and get_expert_linear_names correctly recognize Qwen3_5MoeSparseMoeBlock, but get_experts_list (lines 91–99) dispatches on model type strings extracted as type(model).__name__.lower(). A Qwen3.5 model class like Qwen3_5MoeForCausalLM produces "qwen3_5moeforcausallm", which is not matched in the function's checks. If quantization uses AWQ or NVFP4_SVDQUANT, the code at line 297 in unified_export_hf.py executes get_experts_list(module, model_type) for any MoE module detected by is_moe, triggering NotImplementedError at line 102.

Add "qwen3_5moeforcausallm" to the model type checks in get_experts_list.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/torch/export/layer_utils.py` around lines 328 - 344,
get_experts_list currently dispatches on type(model).__name__.lower() and lacks
handling for Qwen3.5 class names, so add "qwen3_5moeforcausallm" to the
model-type checks inside get_experts_list to match the same Qwen3_5 detection
used by is_moe and get_expert_linear_names; update the conditional branches that
compare model_type (from type(model).__name__.lower()) to include
"qwen3_5moeforcausallm" so get_experts_list returns the correct expert list
instead of raising NotImplementedError (also verify any other qwen3_5 variants
present in that same dispatch and add them if missing).